from flask import request, jsonify, Response, render_template
import requests

import logging
from chroma_utils import getCollection

# 这个函数将在app.py中被调用，用于注册路由
def register_routes(app):
    """
    注册额外的路由到Flask应用
    
    Args:
        app: Flask应用实例
        id: 用于获取ChromaDB集合的ID
    """
    
    @app.route('/webapi/<path:path>', methods=['GET'])
    def proxy(path):
        target_url = f'http://localhost:3000/webapi/{path}'
        app.logger.info(f'代理请求: {target_url}')
        
        try:
            resp = requests.get(
                target_url,
                params=request.args,
                headers={key: value for key, value in request.headers if key != 'Host'},
                stream=True
            )
            app.logger.info(f'代理响应状态码: {resp.status_code}')
            
            response = Response(
                resp.raw.read(),
                status=resp.status_code,
                content_type=resp.headers.get('content-type', 'text/plain')
            )
            
            for key, value in resp.headers.items():
                if key.lower() not in ('content-length', 'transfer-encoding', 'content-encoding'):
                    response.headers[key] = value
            
            return response
            
        except requests.RequestException as e:
            app.logger.error(f'代理请求异常: {str(e)}')
            return jsonify({'error': str(e)}), 500

    @app.route('/chroma_stats')
    def chroma_stats():
        id = request.args.get('id')
        if not id:
            return jsonify({'error': '缺少必要的id参数'}), 400
        collection = getCollection(id)
        document_count = collection.count()
        result = collection.get(limit=1)
        embedding_dimension = len(result['embeddings'][0]) if result and result['embeddings'] else 0
        result = collection.get(limit=5)  # 获取前5个示例文档
        example_docs = result['documents']
        example_ids = result['ids']
        return jsonify({
            'document_count': document_count,
            'embedding_dimension': embedding_dimension,
            'example_documents': example_docs,
            'example_ids': example_ids
        })

    @app.route('/status')
    def status():
        id = request.args.get('id')
        if not id:
            return jsonify({'error': '缺少必要的id参数'}), 400
        collection = getCollection(id)
        all_docs = collection.get()
        return render_template('status.html', documents=all_docs['documents'], ids=all_docs['ids'])


    @app.route('/deletedoc/<file_id>', methods=['DELETE'])
    def delete_document(file_id):
        # 在ChromaDB中删除与该文档相关的数据
        success = chroma_utils.delete_document(file_id)
        if success:
            return jsonify({'success': True, 'message': '文件已删除'})
        else:
            return jsonify({'success': False, 'error': '删除失败'}), 400
